/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.sdk.extensions.sql.impl.transform.agg;



import com.bff.gaia.unified.sdk.annotations.Internal;

import com.bff.gaia.unified.sdk.coders.Coder;

import com.bff.gaia.unified.sdk.coders.CoderRegistry;

import com.bff.gaia.unified.sdk.coders.SerializableCoder;

import com.bff.gaia.unified.sdk.extensions.sql.impl.utils.BigDecimalConverter;

import com.bff.gaia.unified.sdk.schemas.Schema;

import com.bff.gaia.unified.sdk.transforms.Combine;

import com.bff.gaia.unified.sdk.transforms.SerializableFunction;

import com.bff.gaia.unified.sdk.values.Row;

import org.apache.calcite.runtime.SqlFunctions;



import java.math.BigDecimal;

import java.math.MathContext;

import java.math.RoundingMode;

import java.util.stream.StreamSupport;



import static com.bff.gaia.unified.vendor.guava.com.google.common.base.Preconditions.checkArgument;



/**

 * {@link Combine.CombineFn} for <em>Covariance</em> on {@link Number} types.

 *

 * <p>Calculates Population Covariance and Sample Covariance using incremental formulas described in

 * http://en.wikipedia.org/wiki/Algorithms_for_calculating_variance, presumably by Pébay, Philippe

 * (2008), in "Formulas for Robust, One-Pass Parallel Computation of Covariances and Arbitrary-Order

 * Statistical Moments".

 */

@Internal

public class CovarianceFn<T extends Number>

    extends Combine.CombineFn<Row, CovarianceAccumulator, T> {



  static final MathContext MATH_CTX = new MathContext(10, RoundingMode.HALF_UP);



  private static final boolean SAMPLE = true;

  private static final boolean POP = false;



  private boolean isSample; // flag to determine return value should be Covariance Pop or Sample

  private SerializableFunction<BigDecimal, T> decimalConverter;



  public static <V extends Number> CovarianceFn newPopulation(Schema.TypeName typeName) {

    return newPopulation(BigDecimalConverter.forSqlType(typeName));

  }



  public static <V extends Number> CovarianceFn newPopulation(

      SerializableFunction<BigDecimal, V> decimalConverter) {



    return new CovarianceFn<>(POP, decimalConverter);

  }



  public static <V extends Number> CovarianceFn newSample(Schema.TypeName typeName) {

    return newSample(BigDecimalConverter.forSqlType(typeName));

  }



  public static <V extends Number> CovarianceFn newSample(

      SerializableFunction<BigDecimal, V> decimalConverter) {



    return new CovarianceFn<>(SAMPLE, decimalConverter);

  }



  private CovarianceFn(boolean isSample, SerializableFunction<BigDecimal, T> decimalConverter) {

    this.isSample = isSample;

    this.decimalConverter = decimalConverter;

  }



  @Override

  public CovarianceAccumulator createAccumulator() {

    return CovarianceAccumulator.ofZeroElements();

  }



  @Override

  public CovarianceAccumulator addInput(CovarianceAccumulator currentVariance, Row rawInput) {

    if (rawInput == null) {

      return currentVariance;

    }



    checkArgument(rawInput.getFieldCount() > 1);



    return currentVariance.combineWith(

        CovarianceAccumulator.ofSingleElement(

            SqlFunctions.toBigDecimal((Object) rawInput.getValue(0)),

            SqlFunctions.toBigDecimal((Object) rawInput.getValue(1))));

  }



  @Override

  public CovarianceAccumulator mergeAccumulators(Iterable<CovarianceAccumulator> covariances) {

    return StreamSupport.stream(covariances.spliterator(), false)

        .reduce(CovarianceAccumulator.ofZeroElements(), CovarianceAccumulator::combineWith);

  }



  @Override

  public Coder<CovarianceAccumulator> getAccumulatorCoder(

	  CoderRegistry registry, Coder<Row> inputCoder) {

    return SerializableCoder.of(CovarianceAccumulator.class);

  }



  @Override

  public T extractOutput(CovarianceAccumulator accumulator) {

    return decimalConverter.apply(getCovariance(accumulator));

  }



  private BigDecimal getCovariance(CovarianceAccumulator covariance) {



    BigDecimal adjustedCount =

        this.isSample ? covariance.count().subtract(BigDecimal.ONE) : covariance.count();



    return covariance.covariance().divide(adjustedCount, MATH_CTX);

  }

}